import time
import json
import urllib3
import threading
import argparse
from datetime import datetime

class RequestTester:
    def __init__(self, ip_and_port, http_prefix="http"):
        self.http_pool_manager = urllib3.PoolManager()
        self.base_url = f"{http_prefix}://{ip_and_port}/generate_stream"
        self.results = []
        self.running = False
        self.lock = threading.Lock()
        
    def send_request(self, input_content, expected_result="我爱昇腾"):
        """发送单个请求并检查结果"""
        try:
            start_time = time.time()
            
            response = self.http_pool_manager.request(
                "POST",
                self.base_url,
                headers={"Content-Type": "application/json"},
                body=json.dumps({
                    "inputs": input_content,
                    "parameters": {
                        "max_new_tokens": 5,
                        "temperature": 0
                    }
                }).encode(),
                timeout=urllib3.Timeout(total=10)
            )
            
            end_time = time.time()
            response_time = (end_time - start_time) * 1000  # 转换为毫秒
            
            if response.status == 200:
                result_data = json.loads(response.data.decode('utf-8'))
                actual_result = result_data.get("result", "")
                is_match = (actual_result == expected_result)
                
                with self.lock:
                    self.results.append({
                        "timestamp": datetime.now(),
                        "input": input_content,
                        "expected": expected_result,
                        "actual": actual_result,
                        "match": is_match,
                        "response_time": response_time,
                        "status": "success"
                    })
                
                return is_match, actual_result, response_time
            else:
                with self.lock:
                    self.results.append({
                        "timestamp": datetime.now(),
                        "input": input_content,
                        "expected": expected_result,
                        "actual": f"HTTP Error: {response.status}",
                        "match": False,
                        "response_time": response_time,
                        "status": f"error_{response.status}"
                    })
                return False, f"HTTP Error: {response.status}", response_time
                
        except Exception as e:
            error_time = time.time()
            response_time = (error_time - start_time) * 1000 if 'start_time' in locals() else 0
            
            with self.lock:
                self.results.append({
                    "timestamp": datetime.now(),
                    "input": input_content,
                    "expected": expected_result,
                    "actual": f"Exception: {str(e)}",
                    "match": False,
                    "response_time": response_time,
                    "status": "exception"
                })
            return False, f"Exception: {str(e)}", response_time
    
    def start_test(self, input_content, expected_result="我爱昇腾", 
                   requests_per_second=1, duration_seconds=60):
        """开始测试"""
        self.results = []
        self.running = True
        
        print(f"开始测试...")
        print(f"目标地址: {self.base_url}")
        print(f"输入内容: {input_content}")
        print(f"期望结果: {expected_result}")
        print(f"请求频率: {requests_per_second} 次/秒")
        print(f"测试时长: {duration_seconds} 秒")
        print("-" * 50)
        
        start_time = time.time()
        request_count = 0
        
        def send_requests():
            nonlocal request_count
            interval = 1.0 / requests_per_second
            
            while self.running and (time.time() - start_time) < duration_seconds:
                request_start = time.time()
                is_match, actual_result, response_time = self.send_request(
                    input_content, expected_result)
                
                request_count += 1
                status_icon = "✓" if is_match else "✗"
                print(f"[{status_icon}] 请求 {request_count}: "
                      f"实际='{actual_result}', 耗时={response_time:.1f}ms")
                
                # 控制请求频率
                elapsed = time.time() - request_start
                if elapsed < interval:
                    time.sleep(interval - elapsed)
        
        # 启动请求线程
        request_thread = threading.Thread(target=send_requests)
        request_thread.start()
        
        # 等待测试完成
        try:
            while (time.time() - start_time) < duration_seconds and self.running:
                time.sleep(0.1)
        except KeyboardInterrupt:
            print("\n测试被用户中断")
        
        self.running = False
        request_thread.join(timeout=5)
        
        # 生成统计报告
        self.generate_report()
    
    def generate_report(self):
        """生成统计报告"""
        if not self.results:
            print("没有测试结果")
            return
        
        total_requests = len(self.results)
        successful_requests = len([r for r in self.results if r["status"] == "success"])
        matched_requests = len([r for r in self.results if r["match"]])
        
        # 计算响应时间统计
        response_times = [r["response_time"] for r in self.results if r["status"] == "success"]
        avg_response_time = sum(response_times) / len(response_times) if response_times else 0
        max_response_time = max(response_times) if response_times else 0
        min_response_time = min(response_times) if response_times else 0
        
        # 错误分类
        errors = {}
        for result in self.results:
            if not result["match"]:
                error_type = result["status"]
                errors[error_type] = errors.get(error_type, 0) + 1
        
        print("\n" + "="*50)
        print("测试统计报告")
        print("="*50)
        print(f"总请求数: {total_requests}")
        print(f"成功请求: {successful_requests}")
        print(f"匹配期望: {matched_requests}")
        print(f"匹配率: {(matched_requests/total_requests)*100:.2f}%")
        print(f"平均响应时间: {avg_response_time:.2f}ms")
        print(f"最快响应: {min_response_time:.2f}ms")
        print(f"最慢响应: {max_response_time:.2f}ms")
        
        if errors:
            print(f"\n错误统计:")
            for error_type, count in errors.items():
                print(f"  {error_type}: {count} 次")
        
        # 显示一些不匹配的示例
        mismatches = [r for r in self.results if not r["match"] and r["status"] == "success"]
        if mismatches:
            print(f"\n不匹配示例 (前5个):")
            for i, mismatch in enumerate(mismatches[:5]):
                print(f"  {i+1}. 期望='{mismatch['expected']}', 实际='{mismatch['actual']}'")

def main():
    parser = argparse.ArgumentParser(description="API响应测试工具")
    parser.add_argument("--ip", required=True, help="服务器IP和端口，例如 127.0.0.1:8000")
    parser.add_argument("--input", default="测试输入", help="发送的输入内容")
    parser.add_argument("--expected", default="我爱昇腾", help="期望的返回结果")
    parser.add_argument("--frequency", type=float, default=1.0, help="请求频率(次/秒)")
    parser.add_argument("--duration", type=int, default=60, help="测试时长(秒)")
    parser.add_argument("--https", action="store_true", help="使用HTTPS协议")
    
    args = parser.parse_args()
    
    http_prefix = "https" if args.https else "http"
    tester = RequestTester(args.ip, http_prefix)
    
    try:
        tester.start_test(
            input_content=args.input,
            expected_result=args.expected,
            requests_per_second=args.frequency,
            duration_seconds=args.duration
        )
    except KeyboardInterrupt:
        print("\n测试被用户终止")
    except Exception as e:
        print(f"测试出错: {e}")


# USAGE: python tester.py --ip 127.0.0.1:8000 --input "你好" --expected "我爱昇腾" --frequency 2 --duration 30
if __name__ == "__main__":
    main()